import mysql.connector
from mysql.connector import Error
from omegaconf import OmegaConf


# 创建连接
def create_connection():
    connection = None
    try:
        config = OmegaConf.load("application.yaml")
        connection = mysql.connector.connect(
            host=config.mysql.host,
            user=config.mysql.username,
            password=config.mysql.password,
            database=config.mysql.database,
        )
        print("连接成功")
    except Error as e:
        print(f"连接失败: {e}")
    return connection


# 增加数据
def insert_data(connection, query, data):
    cursor = connection.cursor()
    try:
        cursor.execute(query, data)
        connection.commit()  # 提交事务
        print("数据插入成功")
    except Error as e:
        print(f"数据插入失败: {e}")


# 查询数据
def select_data(connection, query):
    cursor = connection.cursor()
    cursor.execute(query)
    result = cursor.fetchall()  # 获取所有结果
    return result


# 删除数据
def delete_data(connection, query, data):
    cursor = connection.cursor()
    try:
        cursor.execute(query, data)
        connection.commit()  # 提交事务
        print("数据删除成功")
    except Error as e:
        print(f"数据删除失败: {e}")


# 更新数据
def update_data(connection, query, data):
    cursor = connection.cursor()
    try:
        cursor.execute(query, data)
        connection.commit()  # 提交事务
        print("数据更新成功")
    except Error as e:
        print(f"数据更新失败: {e}")


def set_log(connection, train_task_id, link):
    """
    设置训练任务日志链接
    :param connection:
    :param train_task_id:
    :param link:
    :return:
    """
    update_query = "UPDATE train_task SET log=%s WHERE id=%s"
    update_data(connection, update_query, (link, train_task_id))


def update_status(connection, train_task_id, status):
    """
    更新训练状态
    :param connection:
    :param train_task_id:
    :param status:
    :return:
    """
    update_query = "UPDATE train_task SET status=%s WHERE id=%s"
    update_data(connection, update_query, (status, train_task_id))


def set_predict_by_id_batch(connection, predicts):
    """
    predicts = [
    ('value1', 1),
    ('value2', 2),
    ('value3', 3),
    ]
    """
    cursor = connection.cursor()
    sql = "UPDATE inference_item SET predict=%s WHERE id=%s"
    cursor.executemany(sql, predicts)
    connection.commit()


def select_mask_color_by_dataset_id(connection, dataset_id):
    select_query = f"SELECT color FROM dataset_category WHERE dataset_id={dataset_id} order by 'order' ASC"
    res = select_data(connection, select_query)
    res = [eval(item[0]) for item in res]
    return res


def check_train_is_pause(connection, train_task_id):
    """判断训练任务是否需要停止"""
    select_query = f"SELECT is_pause FROM train_task WHERE id={train_task_id}"
    res = select_data(connection, select_query)
    return res[0][0]


# 主函数
if __name__ == "__main__":
    connection = create_connection()

    select_query = f"SELECT checkpoint_folder,status FROM train_task WHERE id = {2} limit 1"
    res = select_data(connection, select_query)[0]
    print(res)
